/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.extensions.sql.impl.rel;



import com.bff.gaia.unified.sdk.extensions.sql.impl.transform.agg.AggregationCombineFnAdapter;

import com.bff.gaia.unified.sdk.extensions.sql.impl.utils.CalciteUtils;

import com.bff.gaia.unified.sdk.schemas.Schema;

import com.bff.gaia.unified.sdk.transforms.*;

import com.bff.gaia.unified.sdk.transforms.windowing.*;

import com.bff.gaia.unified.sdk.values.*;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.Lists;

import org.apache.calcite.plan.RelOptCluster;

import org.apache.calcite.plan.RelTraitSet;

import org.apache.calcite.rel.RelNode;

import org.apache.calcite.rel.RelWriter;

import org.apache.calcite.rel.core.Aggregate;

import org.apache.calcite.rel.core.AggregateCall;

import org.apache.calcite.util.ImmutableBitSet;

import com.bff.gaia.unified.sdk.transforms.Combine;

import com.bff.gaia.unified.sdk.transforms.DoFn;

import com.bff.gaia.unified.sdk.transforms.GroupByKey;

import com.bff.gaia.unified.sdk.transforms.PTransform;

import com.bff.gaia.unified.sdk.transforms.ParDo;

import com.bff.gaia.unified.sdk.transforms.WithTimestamps;

import com.bff.gaia.unified.sdk.transforms.windowing.BoundedWindow;

import com.bff.gaia.unified.sdk.transforms.windowing.DefaultTrigger;

import com.bff.gaia.unified.sdk.transforms.windowing.FixedWindows;

import com.bff.gaia.unified.sdk.transforms.windowing.GlobalWindows;

import com.bff.gaia.unified.sdk.transforms.windowing.IntervalWindow;

import com.bff.gaia.unified.sdk.transforms.windowing.Sessions;

import com.bff.gaia.unified.sdk.transforms.windowing.SlidingWindows;

import com.bff.gaia.unified.sdk.transforms.windowing.Window;

import com.bff.gaia.unified.sdk.transforms.windowing.WindowFn;

import com.bff.gaia.unified.sdk.values.KV;

import com.bff.gaia.unified.sdk.values.PCollection;

import com.bff.gaia.unified.sdk.values.PCollectionList;

import com.bff.gaia.unified.sdk.values.Row;

import com.bff.gaia.unified.sdk.values.WindowingStrategy;

import org.joda.time.Duration;



import javax.annotation.Nullable;

import java.io.Serializable;

import java.util.List;



import static java.util.stream.Collectors.toList;

import static com.bff.gaia.unified.vendor.guava.com.google.common.base.Preconditions.checkArgument;



/** {@link UnifiedRelNode} to replace a {@link Aggregate} node. */

public class UnifiedAggregationRel extends Aggregate implements UnifiedRelNode {

  private @Nullable

  WindowFn<Row, IntervalWindow> windowFn;

  private final int windowFieldIndex;



  public UnifiedAggregationRel(

      RelOptCluster cluster,

      RelTraitSet traits,

      RelNode child,

      boolean indicator,

      ImmutableBitSet groupSet,

      List<ImmutableBitSet> groupSets,

      List<AggregateCall> aggCalls,

      @Nullable WindowFn<Row, IntervalWindow> windowFn,

      int windowFieldIndex) {



    super(cluster, traits, child, indicator, groupSet, groupSets, aggCalls);



    this.windowFn = windowFn;

    this.windowFieldIndex = windowFieldIndex;

  }



  @Override

  public RelWriter explainTerms(RelWriter pw) {

    super.explainTerms(pw);

    if (this.windowFn != null) {

      WindowFn windowFn = this.windowFn;

      String window = windowFn.getClass().getSimpleName() + "($" + String.valueOf(windowFieldIndex);

      if (windowFn instanceof FixedWindows) {

        FixedWindows fn = (FixedWindows) windowFn;

        window = window + ", " + fn.getSize().toString() + ", " + fn.getOffset().toString();

      } else if (windowFn instanceof SlidingWindows) {

        SlidingWindows fn = (SlidingWindows) windowFn;

        window =

            window

                + ", "

                + fn.getPeriod().toString()

                + ", "

                + fn.getSize().toString()

                + ", "

                + fn.getOffset().toString();

      } else if (windowFn instanceof Sessions) {

        Sessions fn = (Sessions) windowFn;

        window = window + ", " + fn.getGapDuration().toString();

      } else {

        throw new RuntimeException(

            "Unknown window function " + windowFn.getClass().getSimpleName());

      }

      window = window + ")";

      pw.item("window", window);

    }

    return pw;

  }



  @Override

  public PTransform<PCollectionList<Row>, PCollection<Row>> buildPTransform() {

    Schema outputSchema = CalciteUtils.toSchema(getRowType());

    List<FieldAggregation> aggregationAdapters =

        getNamedAggCalls().stream()

            .map(aggCall -> new FieldAggregation(aggCall.getKey(), aggCall.getValue()))

            .collect(toList());



    return new Transform(

        windowFn, windowFieldIndex, getGroupSet(), aggregationAdapters, outputSchema);

  }



  private static class FieldAggregation implements Serializable {

    final List<Integer> inputs;

    final Combine.CombineFn combineFn;

    final Schema.Field outputField;



    FieldAggregation(AggregateCall call, String alias) {

      inputs = call.getArgList();

      outputField = CalciteUtils.toField(alias, call.getType());

      combineFn =

          AggregationCombineFnAdapter.createCombineFn(

              call, outputField, call.getAggregation().getName());

    }

  }



  private static class Transform extends PTransform<PCollectionList<Row>, PCollection<Row>> {



    private final List<Integer> keyFieldsIds;

    private Schema outputSchema;

    private WindowFn<Row, IntervalWindow> windowFn;

    private int windowFieldIndex;

    private List<FieldAggregation> fieldAggregations;



    private Transform(

        WindowFn<Row, IntervalWindow> windowFn,

        int windowFieldIndex,

        ImmutableBitSet groupSet,

        List<FieldAggregation> fieldAggregations,

        Schema outputSchema) {

      this.windowFn = windowFn;

      this.windowFieldIndex = windowFieldIndex;

      this.fieldAggregations = fieldAggregations;

      this.outputSchema = outputSchema;

      this.keyFieldsIds =

          groupSet.asList().stream().filter(i -> i != windowFieldIndex).collect(toList());

    }



    @Override

    public PCollection<Row> expand(PCollectionList<Row> pinput) {

      checkArgument(

          pinput.size() == 1,

          "Wrong number of inputs for %s: %s",

          UnifiedAggregationRel.class.getSimpleName(),

          pinput);

      PCollection<Row> upstream = pinput.get(0);

      PCollection<Row> windowedStream = upstream;

      if (windowFn != null) {

        windowedStream = assignTimestampsAndWindow(upstream);

      }



      validateWindowIsSupported(windowedStream);



      com.bff.gaia.unified.sdk.schemas.transforms.Group.ByFields<Row> byFields =

          com.bff.gaia.unified.sdk.schemas.transforms.Group.byFieldIds(keyFieldsIds);

      com.bff.gaia.unified.sdk.schemas.transforms.Group.CombineFieldsByFields<Row> combined = null;

      for (FieldAggregation fieldAggregation : fieldAggregations) {

        List<Integer> inputs = fieldAggregation.inputs;

        Combine.CombineFn combineFn = fieldAggregation.combineFn;

        if (inputs.size() > 1 || inputs.isEmpty()) {

          // In this path we extract a Row (an empty row if inputs.isEmpty).

          combined =

              (combined == null)

                  ? byFields.aggregateFieldsById(inputs, combineFn, fieldAggregation.outputField)

                  : combined.aggregateFieldsById(inputs, combineFn, fieldAggregation.outputField);

        } else {

          // Combining over a single field, so extract just that field.

          combined =

              (combined == null)

                  ? byFields.aggregateField(inputs.get(0), combineFn, fieldAggregation.outputField)

                  : combined.aggregateField(inputs.get(0), combineFn, fieldAggregation.outputField);

        }

      }



      PTransform<PCollection<Row>, PCollection<KV<Row, Row>>> combiner = combined;

      if (combiner == null) {

        // If no field aggregations were specified, we run a constant combiner that always returns

        // a single empty row for each key. This is used by the SELECT DISTINCT query plan - in this

        // case a group by is generated to determine unique keys, and a constant null combiner is

        // used.

        combiner = byFields.aggregate(AggregationCombineFnAdapter.createConstantCombineFn());

      }



      return windowedStream

          .apply(combiner)

          .apply("mergeRecord", ParDo.of(mergeRecord(outputSchema, windowFieldIndex)))

          .setRowSchema(outputSchema);

    }



    /** Extract timestamps from the windowFieldIndex, then window into windowFns. */

    private PCollection<Row> assignTimestampsAndWindow(PCollection<Row> upstream) {

      PCollection<Row> windowedStream;

      windowedStream =

          upstream

              .apply(

                  "assignEventTimestamp",

                  WithTimestamps.<Row>of(row -> row.getDateTime(windowFieldIndex).toInstant())

                      .withAllowedTimestampSkew(new Duration(Long.MAX_VALUE)))

              .setCoder(upstream.getCoder())

              .apply(Window.into(windowFn));

      return windowedStream;

    }



    /**

     * Performs the same check as {@link GroupByKey}, provides more context in exception.

     *

     * <p>Verifies that the input PCollection is bounded, or that there is windowing/triggering

     * being used. Without this, the watermark (at end of global window) will never be reached.

     *

     * <p>Throws {@link UnsupportedOperationException} if validation fails.

     */

    private void validateWindowIsSupported(PCollection<Row> upstream) {

      WindowingStrategy<?, ?> windowingStrategy = upstream.getWindowingStrategy();

      if (windowingStrategy.getWindowFn() instanceof GlobalWindows

          && windowingStrategy.getTrigger() instanceof DefaultTrigger

          && upstream.isBounded() != PCollection.IsBounded.BOUNDED) {



        throw new UnsupportedOperationException(

            "Please explicitly specify windowing in SQL query using HOP/TUMBLE/SESSION functions "

                + "(default trigger will be used in this case). "

                + "Unbounded input with global windowing and default trigger is not supported "

                + "in Unified SQL aggregations. "

                + "See GroupByKey section in Unified Programming Guide");

      }

    }



    static DoFn<KV<Row, Row>, Row> mergeRecord(Schema outputSchema, int windowStartFieldIndex) {

      return new DoFn<KV<Row, Row>, Row>() {

        @ProcessElement

        public void processElement(

			@Element KV<Row, Row> kvRow, BoundedWindow window, OutputReceiver<Row> o) {

          List<Object> fieldValues =

              Lists.newArrayListWithCapacity(

                  kvRow.getKey().getValues().size() + kvRow.getValue().getValues().size());



          fieldValues.addAll(kvRow.getKey().getValues());

          fieldValues.addAll(kvRow.getValue().getValues());



          if (windowStartFieldIndex != -1) {

            fieldValues.add(windowStartFieldIndex, ((IntervalWindow) window).start());

          }



          o.output(Row.withSchema(outputSchema).addValues(fieldValues).build());

        }

      };

    }

  }



  @Override

  public Aggregate copy(

      RelTraitSet traitSet,

      RelNode input,

      boolean indicator,

      ImmutableBitSet groupSet,

      List<ImmutableBitSet> groupSets,

      List<AggregateCall> aggCalls) {

    return new UnifiedAggregationRel(

        getCluster(),

        traitSet,

        input,

        indicator,

        groupSet,

        groupSets,

        aggCalls,

        windowFn,

        windowFieldIndex);

  }

}